package com.timeriver.feature_project

import org.apache.spark.ml.feature.QuantileDiscretizer
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * QuantileDiscretizer接收一个具有连续特征的列，并输出一个具有分类特征的列。
  * 桶的数量是由numBuckets参数设置的。
  * 使用的桶的数量有可能小于这个值，例如，如果输入的不同值太少，无法创建足够多的不同量级。
  */
object QuantileDiscretizerDemo {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .master("local[6]")
      .appName("连续数据等分位离散化")
      .getOrCreate()

    val tuples: Array[(Int, Double)] = Array((0, 18), (1, 32), (2, 5), (3, 12), (4, 56), (5, 67), (6, 89))
    val data: DataFrame = session.createDataFrame(tuples).toDF("id", "label")

    val discretizer: QuantileDiscretizer = new QuantileDiscretizer()
      .setInputCol("label")
      .setOutputCol("quantileDiscretizer_label")
      .setNumBuckets(5)

    val res: DataFrame = discretizer.fit(data).transform(data)
    res.show(false)
  }
}
